/*
 * Copyright (C) 2017-2019 Dremio Corporation
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.dremio.sabot.op.sort.topn;

import java.util.concurrent.TimeUnit;

import javax.inject.Named;

import org.apache.arrow.memory.BufferAllocator;
import org.apache.arrow.vector.types.pojo.Schema;

import com.dremio.common.AutoCloseables;
import com.dremio.exec.exception.SchemaChangeException;
import com.dremio.exec.record.VectorAccessible;
import com.dremio.exec.record.VectorContainer;
import com.dremio.exec.record.selection.SelectionVector2;
import com.dremio.exec.record.selection.SelectionVector4;
import com.dremio.sabot.exec.context.FunctionContext;
import com.dremio.sabot.op.sort.external.RecordBatchData;
import com.dremio.sabot.op.sort.external.Sv4HyperContainer;
import com.google.common.base.Stopwatch;
import com.google.common.base.Throwables;

import io.netty.buffer.ArrowBuf;

public abstract class PriorityQueueTemplate implements PriorityQueue {
  private static final org.slf4j.Logger logger = org.slf4j.LoggerFactory.getLogger(PriorityQueueTemplate.class);

  private SelectionVector4 heapSv4; //This holds the heap
  private SelectionVector4 finalSv4; //This is for final sorted output
  private Sv4HyperContainer hyperBatch;
  private FunctionContext context;
  private BufferAllocator allocator;
  private int limit;
  private int queueSize = 0;
  private int batchCount = 0;
  private boolean hasSv2;
  private int maxSize;

  @Override
  public void init(Sv4HyperContainer hyperBatch, int limit, FunctionContext context, BufferAllocator allocator,  boolean hasSv2, int maxSize) {
    this.limit = limit;
    this.context = context;
    this.allocator = allocator;
    final ArrowBuf ArrowBuf = allocator.buffer(4 * (limit + 1));
    this.heapSv4 = new SelectionVector4(ArrowBuf, limit, Character.MAX_VALUE);
    this.hasSv2 = hasSv2;
    this.hyperBatch = hyperBatch;
    this.maxSize = maxSize;
    doSetup(context, hyperBatch, null);
  }

  @Override
  public void resetQueue(final VectorContainer newQueue, final SelectionVector4 oldHeap) {
    Schema schema = hyperBatch.getSchema();

    // Cleanup old data.
    close();

    // generating
    hyperBatch = new Sv4HyperContainer(allocator, schema);
    hyperBatch.addBatch(newQueue);
    batchCount = hyperBatch.iterator().next().getValueVectors().length;
    final ArrowBuf ArrowBuf = allocator.buffer(4 * (limit + 1));
    heapSv4 = new SelectionVector4(ArrowBuf, limit, Character.MAX_VALUE);
    // Reset queue size (most likely to be set to limit).
    queueSize = 0;
    for (int i = 0; i < oldHeap.getTotalCount(); i++) {
      heapSv4.set(i, oldHeap.get(i));
      ++queueSize;
    }
    oldHeap.clear();
    doSetup(context, hyperBatch, null);
  }

  @Override
  public void add(RecordBatchData batch) throws SchemaChangeException{
    Stopwatch watch = Stopwatch.createStarted();

    hyperBatch.addBatch(batch.getContainer());

    doSetup(context, hyperBatch, null);

    int count = 0;
    SelectionVector2 sv2 = null;
    if (hasSv2) {
      sv2 = batch.getSv2();
    }
    for (; queueSize < limit && count < batch.getRecordCount();  count++) {
      heapSv4.set(queueSize, batchCount, hasSv2 ? sv2.getIndex(count) : count);
      queueSize++;
      siftUp();
    }
    for (; count < batch.getRecordCount(); count++) {
      heapSv4.set(limit, batchCount, hasSv2 ? sv2.getIndex(count) : count);
      if (compare(limit, 0) < 0) {
        swap(limit, 0);
        siftDown();
      }
    }
    batchCount++;
    if (hasSv2) {
      sv2.clear();
    }
    logger.debug("Took {} us to add {} records", watch.elapsed(TimeUnit.MICROSECONDS), count);
  }

  @Override
  public void generate() throws SchemaChangeException {
    Stopwatch watch = Stopwatch.createStarted();
    final ArrowBuf ArrowBuf = allocator.buffer(4 * queueSize);
    finalSv4 = new SelectionVector4(ArrowBuf, queueSize, maxSize);
    for (int i = queueSize - 1; i >= 0; i--) {
      finalSv4.set(i, pop());
    }
    logger.debug("Took {} us to generate output of {}", watch.elapsed(TimeUnit.MICROSECONDS), finalSv4.getTotalCount());
  }

  @Override
  public Sv4HyperContainer getHyperBatch() {
    return hyperBatch;
  }

  @Override
  public SelectionVector4 getHeapSv4() {
    return heapSv4;
  }

  @Override
  public SelectionVector4 getFinalSv4() {
    return finalSv4;
  }

  @Override
  public void close() {
    try{
      AutoCloseables.close(heapSv4, hyperBatch, finalSv4);
    }catch(Exception ex){
      throw Throwables.propagate(ex);
    }
  }

  private void siftUp() {
    int p = queueSize - 1;
    while (p > 0) {
      if (compare(p, (p - 1) / 2) > 0) {
        swap(p, (p - 1) / 2);
        p = (p - 1) / 2;
      } else {
        break;
      }
    }
  }

  private void siftDown() {
    int p = 0;
    int next;
    while (p * 2 + 1 < queueSize) {
      if (p * 2 + 2 >= queueSize) {
        next = p * 2 + 1;
      } else {
        if (compare(p * 2 + 1, p * 2 + 2) >= 0) {
          next = p * 2 + 1;
        } else {
          next = p * 2 + 2;
        }
      }
      if (compare(p, next) < 0) {
        swap(p, next);
        p = next;
      } else {
        break;
      }
    }
  }

  public int pop() {
    int value = heapSv4.get(0);
    swap(0, queueSize - 1);
    queueSize--;
    siftDown();
    return value;
  }

  public void swap(int sv0, int sv1) {
    int tmp = heapSv4.get(sv0);
    heapSv4.set(sv0, heapSv4.get(sv1));
    heapSv4.set(sv1, tmp);
  }

  public int compare(int leftIndex, int rightIndex) {
    int sv1 = heapSv4.get(leftIndex);
    int sv2 = heapSv4.get(rightIndex);
    return doEval(sv1, sv2);
  }

  public abstract void doSetup(@Named("context") FunctionContext context, @Named("incoming") VectorAccessible incoming, @Named("outgoing") VectorAccessible outgoing);
  public abstract int doEval(@Named("leftIndex") int leftIndex, @Named("rightIndex") int rightIndex);

}
